{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "from dataclasses import dataclass\n",
    "from glob import glob\n",
    "from itertools import chain\n",
    "from pathlib import Path\n",
    "from typing import Any, Dict, List, Optional, Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from PIL import Image\n",
    "from torch.utils.data import Dataset\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import AutoProcessor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37e85453cf744134ab03088c311c6d86",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2290793 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b286ef9c519c40bfb55d34583d3006f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "writing:   0%|          | 0/2290793 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image_dir = \"data/jackyhate/unzip2mdata\"\n",
    "\n",
    "\n",
    "def laod_data_from_dir(dir: str):\n",
    "    image_text_dir = Path(dir)\n",
    "    # image_file_list = list(image_text_dir.glob(\"*.jpg\"))\n",
    "    text_file_list = list(image_text_dir.glob(\"*/*.json\"))#[:100]\n",
    "\n",
    "    def load_text(file_path: Path) -> Dict[str, Any]:\n",
    "        with open(file_path, \"r\") as f:\n",
    "            return json.load(f).get(\"prompt\")\n",
    "\n",
    "    text_data_list = [\n",
    "        {\n",
    "            \"prompt\": load_text(file_path),\n",
    "            \"file_path\": str(file_path)#.replace(\".json\", \"\"),\n",
    "        }\n",
    "        for file_path in tqdm(text_file_list)\n",
    "    ]\n",
    "    return text_data_list\n",
    "\n",
    "\n",
    "final_data = laod_data_from_dir(image_dir)\n",
    "\n",
    "with open(\"data/jackyhate/text-to-image-2M/data_2m_p1_large.json\", \"w\") as f:\n",
    "    for i in tqdm(final_data,desc='writing'):\n",
    "        f.write(json.dumps(i, ensure_ascii=False) + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'prompt': \"A diagram illustrating various microprocessor vendors' debugging extensions, including Infineon, MIPS, and Freescale, with a focus on JTAG support.\",\n",
       "  'file_path': PosixPath('data/jackyhate/unzip2mdata/data_000030/llava_next_captioned_700K_00024770.json')}]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# final_data[:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Image.open(\"data/jackyhate/unzip2mdata/data_000030/llava_next_captioned_700K_00024770.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
